import os
import uuid
import json
import time
import requests
import cv2
import urllib3
from typing import Optional, Tuple, Dict, Any
from dotenv import load_dotenv

from utils import sr_s3
from utils.sr_utils import video_compress_image_size_path


class MiniMaxVideoGenerator:
    """MiniMax视频生成API封装类（适配modelType和模型字典）"""

    def __init__(self):
        load_dotenv(verbose=True, override=True)
        self.MINIMAX_URL = os.getenv('MINIMAX_URL', 'https://api.minimax.io')
        self.API_KEY = os.getenv('MINIMAX_API_KEY')

        # 模型字典：编号 -> 模型名称（按需求定义）
        self.MODEL_DICT = {
            "MiniMax-Hailuo-02",
            "T2V-01-Director",
            "I2V-01-Director",
            "I2V-01-live",
            "MiniMax-Hailuo-02",
        }

        # 模型对应的任务类型（文生/图生）
        self.MODEL_TASK_MAP = {
            1: "text",  # MiniMax-Hailuo-02 支持文生/图生，这里默认文生
            2: "text",  # T2V-01-Director 文生
            3: "image",  # I2V-01-Director 图生
            4: "image",  # I2V-01-live 图生
            5: "image",  # MiniMax-Hailuo-02 图生
        }

        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.API_KEY}"
        }

        self.temp_image_dir = "./upload"
        self.temp_video_dir = "./download"
        os.makedirs(self.temp_image_dir, exist_ok=True)
        os.makedirs(self.temp_video_dir, exist_ok=True)
        self.verify_ssl = True

    # --------------------------
    # 核心适配：参数校验逻辑
    # --------------------------
    def _validate_model_and_type(self, model_id: int, model_type: int) -> Tuple[bool, str]:
        """验证模型编号与modelType是否匹配（1=文生，2=图生）"""
        if model_id not in self.MODEL_DICT:
            return False, f"不支持的模型{model_id}"

        # 模型对应的任务类型（内部映射）
        inner_task_type = self.MODEL_TASK_MAP[model_id]
        print(inner_task_type)

        # # 校验modelType（1=文生，2=图生）
        if model_type == 1 and inner_task_type not in ("text", "image"):
            # modelType=1（文生）仅支持文生模型
            return False, f"模型{self.MODEL_DICT[model_id]}不支持文生视频（modelType=1）"
        if model_type == 2 and inner_task_type not in ("image", "subject"):
            # modelType=2（图生）仅支持图生模型
            return False, f"模型{self.MODEL_DICT[model_id]}不支持图生视频（modelType=2）"

        return True, ""

    def _validate_model_params(self, model_id: int, params: Dict[str, Any]) -> Tuple[bool, str]:
        """验证模型专属参数"""
        model_name = self.MODEL_DICT[model_id]

        # 1. 文生模型参数校验（T2V-01-Director、MiniMax-Hailuo-02）
        if model_id in (1, 2):
            # 文生模型不需要首帧图片，但需要prompt
            if not params.get("prompt"):
                return False, f"{model_name}（文生）需要prompt参数"
            return True, ""

        # 2. 图生模型参数校验（I2V-01-Director、I2V-01-live）
        if model_id in (3, 4, 5):
            # 必须提供首帧图片
            if not params.get("first_image"):
                return False, f"{model_name}（图生）需要first_image参数"
            return True, ""

        # # 3. 主体参考模型校验（S2V-01）
        # if model_id == 6:
        #     if not params.get("subject_reference"):
        #         return False, f"{model_name}需要subject_reference参数（主体图片）"
        #     return True, ""
        #
        # return True, ""

    def _validate_duration_resolution(self, model_id: int, duration: int, resolution: str) -> Tuple[bool, str]:
        """验证时长和分辨率（按模型限制）"""
        model_name = self.MODEL_DICT[model_id]

        # 01系列模型（T2V-01-Director、I2V-01-Director、I2V-01-live、S2V-01）
        if model_id in (2, 3, 4):
            if duration != 6:
                return False, f"{model_name}（01系列）仅支持6秒时长"
            if resolution is not None:
                return False, f"{model_name}（01系列）不支持自定义分辨率（默认720P）"

        # MiniMax-Hailuo-02（02系列）
        if model_id == 1 or model_id == 5:
            if duration not in (6, 10):
                return False, f"{model_name}支持6秒或10秒时长"
            if resolution == "1080P" and duration == 10:
                return False, f"{model_name}的1080P分辨率仅支持6秒时长"
            if resolution and resolution not in ("768P", "1080P"):
                return False, f"{model_name}支持768P或1080P分辨率"

        return True, ""

    # --------------------------
    # 工具方法：文件处理（复用）
    # --------------------------
    def _delete_temp_files(self, *file_paths: str) -> None:
        for path in file_paths:
            if path and os.path.exists(path):
                try:
                    os.remove(path)
                except Exception as e:
                    print(f"清理临时文件失败 {path}: {str(e)}")

    def _download_image(self, img_url: str, storage_service: str, bucket_name: str) -> Tuple[bool, str]:
        try:
            name = str(uuid.uuid4())
            ext = os.path.splitext(img_url)[-1].lower() or '.jpg'
            save_path = os.path.join(self.temp_image_dir, f"{name}{ext}")

            from utils import sr_s3
            success, err = sr_s3.download_file(
                img_url, save_path, serviceName=storage_service, bucketName=bucket_name
            )
            if not success:
                return False, f"图片下载失败：{err}"

            success, new_path = video_compress_image_size_path(save_path, 10, (300, 6000), (300, 6000))
            if not success:
                return False, "图片压缩失败"
            # if ext not in ('.jpg', '.jpeg', '.png'):
            #     return False, "图片格式仅支持jpg/png"
            # if os.path.getsize(save_path) > 20 * 1024 * 1024:
            #     return False, "图片大小超过20MB限制"
            return True, new_path
        except Exception as e:
            return False, f"图片处理异常：{str(e)}"

    def _image_to_base64(self, img_path: str) -> Tuple[bool, str]:
        try:
            import base64
            with open(img_path, "rb") as f:
                base64_str = base64.b64encode(f.read()).decode("utf-8")
            return True, f"data:image/jpeg;base64,{base64_str}"
        except Exception as e:
            return False, f"图片转base64失败：{str(e)}"

    # --------------------------
    # 核心方法：API调用逻辑
    # --------------------------
    def _create_video_task(self, payload: dict) -> Tuple[bool, str, str]:
        try:
            url = f"{self.MINIMAX_URL}/v1/video_generation"
            print(url)
            response = requests.post(
                url=url,
                headers=self.headers,
                data=json.dumps(payload),
                verify=self.verify_ssl
            )
            response_data = response.json()
            print(response_data)
            base_resp = response_data.get("base_resp", {})

            if base_resp.get("status_code") == 0:
                task_id = response_data.get("task_id")
                return (True, task_id, "") if task_id else (False, "", "无task_id")
            else:
                err_msg = base_resp.get("status_msg", "未知错误")
                return False, "", f"创建失败：{err_msg}（code: {base_resp.get('status_code')}）"
        except Exception as e:
            return False, "", f"请求异常：{str(e)}"

    def _poll_task_status(self, task_id: str, timeout: int = 300, interval: int = 10) -> tuple[bool, str] | tuple[bool, str, Any] | tuple[tuple[bool, str], int]:
        if not task_id:
            return False, "任务ID为空"

        url = f"{self.MINIMAX_URL}/v1/query/video_generation"
        start_time = time.time()

        while time.time() - start_time < timeout:
            try:
                params = {"task_id": task_id}
                response = requests.get(
                    url,
                    headers=self.headers,
                    params=params,
                    verify=self.verify_ssl,
                    timeout=(5, 15)
                )
                response_data = response.json()
                print(response_data)
                base_resp = response_data.get("base_resp", {})

                if base_resp.get("status_code") != 0:
                    err_msg = base_resp.get("status_msg", "查询失败")
                    return False, f"查询失败：{err_msg}", base_resp.get("status_code")

                status = response_data.get("status")
                if status == "Success":
                    file_id = response_data.get("file_id")
                    success, video_url = self._get_video_url(file_id)
                    return success, video_url, status
                elif status in ("Fail", "Unknown"):
                    return False, f"任务失败：{status}", 0
                else:
                    print(f"任务状态：{status}，等待{interval}秒...")
                    time.sleep(interval)

            except Exception as e:
                return False, f"轮询异常：{str(e)}", 0

        return False, f"超时（{timeout}秒）", 0

    def _get_video_url(self, file_id: str) -> Tuple[bool, str]:
        try:
            url = f"{self.MINIMAX_URL}/v1/files/retrieve"
            params = {"file_id": file_id}
            response = requests.get(
                url,
                headers=self.headers,
                params=params,
                verify=self.verify_ssl
            )
            response_data = response.json()
            if response_data.get("base_resp", {}).get("status_code") == 0:
                return True, response_data.get("file", {}).get("download_url", "")
            else:
                return False, "获取下载链接失败"
        except Exception as e:
            return False, f"获取视频URL异常：{str(e)}"

    # --------------------------
    # 外部调用入口（核心逻辑）
    # --------------------------
    def process_video_task(self, task_params: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理视频生成任务（外部调用入口）

        参数:
            task_params: 包含以下关键参数：
                - taskId: 自定义任务ID
                - modelId: 模型编号（1-5，对应MODEL_DICT）
                - modelType: 1=文生视频，2=图生视频
                - prompt: 提示词（文生必填）
                - duration: 视频时长（默认6秒）
                - resolution: 分辨率（部分模型支持）
                - first_image: 首帧图片URL（图生必填）
                - subject_reference: 主体图片URL（S2V-01必填）
                - storageService/bucketName/storedPrefix/QueueResult: 存储和队列参数
        """
        task_id = task_params.get('taskId', f"minimax_{uuid.uuid4().hex[:8]}")
        temp_files = []

        try:
            # 1. 基础参数校验
            required_params = ['moduleName', 'modelType', 'storageService', 'bucketName', 'storedPrefix', 'QueueResult']
            for param in required_params:
                if not task_params.get(param):
                    return {"status": 2, "msg": f"缺少必填参数：{param}", "code": 6003}

            # 2. 解析核心参数
            model_name = task_params['moduleName']
            model_type = int(task_params['modelType'])  # 1=文生，2=图生

            prompt = task_params.get('prompt', '')
            duration = int(task_params.get('duration', 6))
            resolution = task_params.get('mResolution')
            if resolution:
                resolution = resolution.upper()
            print(resolution)
            # # 3. 模型与类型匹配校验
            # valid, err = self._validate_model_and_type(model_name, model_type)
            # if not valid:
            #     return {"status": 2, "msg": err}
            #
            # # 4. 模型专属参数校验
            # valid, err = self._validate_model_params(model_name, task_params)
            # if not valid:
            #     return {"status": 2, "msg": err}
            #
            # # 5. 时长与分辨率校验
            # valid, err = self._validate_duration_resolution(model_name, duration, resolution)
            # if not valid:
            #     return {"status": 2, "msg": err}

            # 6. 构建payload（按模型定制）
            payload = {
                "model": model_name,
                "prompt": prompt,
                "duration": duration,
                "prompt_optimizer": task_params.get('prompt_optimizer', True)
            }

            # 7. 按模型添加专属参数
            if model_type == 2:  # 图生模型（I2V系列）
                # 处理首帧图片
                payload["first_frame_image"] = sr_s3.get_url(task_params.get('first_image'), serviceName=task_params.get('storageService', "s3"), bucketName=task_params.get('bucketName', ""))
            if model_type == 3:  # 主体参考模型（S2V-01）
                # 处理主体图片
                success, subj_path = self._download_image(
                    task_params['subject_reference'],
                    task_params['storageService'],
                    task_params['bucketName']
                )
                if not success:
                    return {"status": 0, "msg": subj_path, "code": 6005}
                temp_files.append(subj_path)

                success, subj_base64 = self._image_to_base64(subj_path)
                if not success:
                    return {"status": 0, "msg": subj_base64, "code": 6005}
                payload["subject_reference"] = [subj_base64]  # 数组格式

            if model_name == "MiniMax-Hailuo-02" and resolution:  # MiniMax-Hailuo-02支持分辨率
                payload["resolution"] = resolution

            # 8. 创建任务
            success, task_api_id, err_msg = self._create_video_task(payload)
            if not success:
                return {"status": 0, "msg": err_msg, "code": 6001}
            print(f"任务{task_id}提交成功，API任务ID：{task_api_id}")
            # task_api_id = "299295099580539"
            # 9. 轮询结果
            success, video_url, code = self._poll_task_status(
                task_api_id,
                timeout=task_params.get('timeout', 300),
                interval=task_params.get('interval', 10)
            )

            if not success:
                if code == 1026 or code == 1027:
                    return {"status": 7, "msg": f"任务处理失败：{video_url}", "code": 7202}
                elif code == 1033:
                    return {"status": 0, "msg": f"第三方系统错误：{video_url}", "code": 6001}
                return {"status": 0, "msg": f"任务处理失败：{video_url}", "code": 6001}

            # 10. 下载并上传视频
            success, video_info = self._download_video(video_url, task_id)
            if not success:
                return {"status": 0, "msg": video_info.get("error")}
            # temp_files.append(video_info["save_path"])
            success, s3_key = self._upload_video_to_s3(
                video_info["save_path"],
                task_params['storedPrefix'],
                task_params['storageService'],
                task_params['bucketName']
            )
            if not success:
                return {"status": 0, "msg": f"上传S3失败：{s3_key}", "code": 6001}

            # 11. 返回结果
            msg_body = {
                'status': 1,
                'width': video_info['width'],
                'height': video_info['height'],
                'local_path': video_info["save_path"],
                'url': s3_key,
                "model_mode": "std",  # 模型类型标准模式
                'msg': "success",
                "code": ""
            }
            return msg_body
        except Exception as e:
            return {"status": 0, "msg": f"处理异常：{str(e)}", "code": 6001}
        #
        # finally:
        #     self._delete_temp_files(*temp_files)

    # --------------------------
    # 复用方法：下载与上传
    # --------------------------
    def _download_video(self, video_url: str, task_id: str) -> Tuple[bool, Dict[str, Any]]:
        try:
            filename = f"minimax_video_{task_id}.mp4"
            save_path = os.path.join(self.temp_video_dir, filename)

            response = requests.get(
                video_url,
                stream=True,
                timeout=30,
                verify=self.verify_ssl
            )
            response.raise_for_status()

            with open(save_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        f.write(chunk)

            cap = cv2.VideoCapture(save_path)
            video_info = {
                "save_path": save_path,
                "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                "duration": cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 0
            }
            cap.release()
            return True, video_info
        except Exception as e:
            return False, {"error": f"视频下载失败：{str(e)}"}

    def _upload_video_to_s3(self, video_path: str, s3_storage: str, storage_service: str, bucket_name: str) -> Tuple[bool, str]:
        if not video_path or not os.path.exists(video_path):
            return False, "视频文件不存在"
        try:
            from utils import sr_s3
            s3_key = sr_s3.upload_login_or_visitor(
                video_path, s3_storage, serviceName=storage_service, bucketName=bucket_name
            )
            return True, s3_key
        except Exception as e:
            return False, f"上传S3失败：{str(e)}"


# --------------------------
# 调用示例
# --------------------------
if __name__ == "__main__":
    generator = MiniMaxVideoGenerator()

    # 示例1：文生视频（modelType=1，模型编号2=T2V-01-Director）
    # text_task = {
    #     "taskId": "text_task_001",
    #     "modelId": 2,  # 对应T2V-01-Director
    #     "modelType": 1,  # 1=文生视频
    #     "prompt": "一只猫在草地上玩球[推进]",
    #     "duration": 6,
    #     "storageService": "r2",
    #     "bucketName": "picwand",
    #     "storedPrefix": "video/visitor/minimax",
    #     "QueueResult": "AILocalVideoResult"
    # }

    # 示例2：图生视频（modelType=2，模型编号3=I2V-01-Director）
    image_task = {'taskId': '0235185679cdba8fb296', 'prompt': '有个人把橘子吃掉了', 'first_image': 'video/image_to_video/upload/d41d8cd98f00b204e9800998ecf8427e/1755656762368.jpg', 'last_image': '', 'duration': '6', 'fps': '25', 'resolution': '720p', 'ratio': '', 'appId': '1', 'taskType': 'pwai_video_itf', 'QueueResult': 'AILocalVideoResult', 'storageService': 'r2', 'isLogin': 1, 'bucketName': 'picwand', 'storedPrefix': 'video/image_to_video', 'progress': 1, 'queueName': 'AIHubVideoTasks_Test', 'modelType': 2, 'moduleType': 25, 'moduleName': 'I2V-01-live', 'mResolution': '720p', 'mode': 'std', 'negative_prompt': '', 'cfg_scale': 0.5, 'p_t_b': 0, 'n_t_b': 0, 'motion_type': '', 'model_style': '', 'camera_movement': ''}

    # 执行任务
    result = generator.process_video_task(image_task)  # 切换为image_task测试图生
    print(result)
